package siberia.lexer;

import java.util.HashMap;
import java.util.Map;
import java.util.ArrayList;

import siberia.*;

public class Lexer {
    private static Map<String, TokenType> keywords = new HashMap<String, TokenType>() {
        {
            put("true", TokenType.True);
            put("false", TokenType.False);
            put("func", TokenType.Func);
            put("if", TokenType.If);
            put("elif", TokenType.Elif);
            put("else", TokenType.Else);
            put("for", TokenType.For);
            put("while", TokenType.While);
            put("return", TokenType.Return);
            put("int", TokenType.Int);
            put("bool", TokenType.Bool);
            put("string", TokenType.Str);
            put("float", TokenType.Float);
        }
    };

    private String source;
    private int pos = 0;
    private int line = 1;
    private int start = 0;
    private ArrayList<Token> tokens = new ArrayList<>();

    public Lexer(String src) {
        source = src;
    }

    public ArrayList<Token> scanTokens() {
        while (!isAtEnd()) {
            start = pos;
            scanToken();
        }

        tokens.add(new Token(TokenType.EndOfFile, null, line));
        return tokens;
    }

    private void scanToken() {
        char c = advance();
        if (c == '(') addToken(TokenType.LParen);  
        else if (c == ')') addToken(TokenType.RParen);  
        else if (c == '[') addToken(TokenType.LBracket);  
        else if (c == ']') addToken(TokenType.RBracket);  
        else if (c == '{') addToken(TokenType.LBrace);  
        else if (c == '}') addToken(TokenType.RBrace);  
        else if (c == ',') addToken(TokenType.Comma);  
        else if (c == '-') addToken(TokenType.Minus);  
        else if (c == '+') addToken(TokenType.Add);  
        else if (c == ';') addToken(TokenType.Semicolon);  
        else if (c == '*') addToken(TokenType.Mul);  
        else if (c == '%') addToken(TokenType.Mod);  
        else if (c == '!') addToken(match('=') ? TokenType.NotEqual : TokenType.Not);
        else if (c == '<') addToken(match('=') ? TokenType.LessEqual : TokenType.LessThan);
        else if (c == '>') addToken(match('=') ? TokenType.GreaterEqual : TokenType.GreaterThan);
        else if (c == '=') addToken(match('=') ? TokenType.Equal : TokenType.Assign);
        else if (c == '/') addToken(match('/') ? TokenType.RoundDiv : TokenType.Div);
        else if (c == '\'' || c == '"') string(c);
        else if (c == ' ' || c == '\r' || c == '\t');
        else if (c == '\n') {
            line++;
        } else if (c == '#') {
            if (match('=')) skipMultilineComment();
            else skipComment();
        } else {
            if (isDigit(c)) {
                number();
            } else if (isAlpha(c)) {
                identifier();
            } else {
                Siberia.error(line, "Unexpected character.");
            }
        }
    }

    private void skipComment() {
        while (peek() != '\n' && !isAtEnd()) advance();
    }

    private void skipMultilineComment() {
        while (peek() != '=' || peekNext() != '#' && !isAtEnd()) {
            if (peek() != '\n') advance();
            else {
                line++;
                advance();
            }
        }
        advance();
        advance();
    }

    private void identifier() {
        while (isAlphaNumeric(peek())) advance();

        String text = source.substring(start, pos);
        TokenType type = keywords.get(text);
        if (type == null) type = TokenType.Identifier;
        addToken(type, text);
    }

    private void number() {
        while (isDigit(peek())) advance();

        boolean isInteger = true;

        if (peek() == '.' && isDigit(peekNext())) {
            isInteger = false;
            advance();
            
            while (isDigit(peek())) advance();
        }
        if (isInteger) {
            addToken(TokenType.IntLiteral, Integer.parseInt(source.substring(start, pos)));
        } else {
            addToken(TokenType.FloatLiteral, Double.parseDouble(source.substring(start, pos)));
        }
    }

    private void string(char endChar) {
        while (peek() != endChar && !isAtEnd()) {
            if (peek() == '\n') line++;
            advance();
        }

        if (isAtEnd()) {
            Siberia.error(line, "Unterminated string.");
            return;
        }

        advance();

        String value = source.substring(start + 1, pos - 1);
        addToken(TokenType.StringLiteral, value);
    }

    private boolean match(char expected) {
        if (isAtEnd()) return false;
        if (source.charAt(pos) != expected) return false;

        pos++;
        return true;
    }

    private char peek() {
        if (isAtEnd()) return '\0';
        return source.charAt(pos);
    }

    private char peekNext() {
        if (pos + 1 >= source.length()) return '\0';
        return source.charAt(pos + 1);
    }

    private boolean isAlpha(char c) {
        return (c >= 'a' && c <= 'z') ||
                (c >= 'A' && c <= 'Z') ||
                (c == '_');
    }

    private boolean isAlphaNumeric(char c) {
        return isAlpha(c) || isDigit(c);
    }

    private boolean isDigit(char c) {
        return c >= '0' && c <= '9';
    }

    private boolean isAtEnd() {
        return pos >= source.length();
    }

    private char advance() {
        return source.charAt(pos++);
    }

    private void addToken(TokenType type) {
        addToken(type, null);
    }

    private void addToken(TokenType type, Object literal) {
        String text = source.substring(start, pos);
        Token t = new Token(type, literal, line);
        tokens.add(t);
    }
}
